import matplotlib
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
import pandas as pd
from resultAnalysis import ResultAnalysis
import numpy as np


class Visualization:
    initial_money = 100000  # 等待数据传入

    def __init__(self, initial):
        self.initial_money = initial

    def plot_return(self, return_list, isSSEC=False, isCSI300=False, start_date='2010-01-04', end_date='2019-12-31', strategy='策略'):
        '''
        绘制资产变化曲线，投资开始和结束时间可以自由设定，可以不是交易日。
        :param return_list: 策略资产变化曲线
        :param isSSEC: 是否绘制上证指数
        :param isCSI300: 是否绘制沪深300指数
        :param start_date: 策略开始时间
        :param end_date: 策略结束时间
        '''

        if isSSEC and not isCSI300:
            ra = ResultAnalysis(initial=self.initial_money, start_date=start_date, end_date=end_date)
            days, return_list_SSEC = ra.cal_return_list_SSEC()
            # num_days = len(return_list_SSEC)
            # days = list(range(num_days))
            # print(len(days))
            fig, ax = plt.figure()
            ax.ticklabel_format(style='plain')
            # 根据情况画一个初始资金虚线
            plt.plot(days, return_list, color='red')
            plt.plot(days, return_list_SSEC, color='green')
            plt.xlabel('时间')
            plt.ylabel('总资产（元）')
            plt.axhline(ra.get_initial_money(), linestyle='--', color='black')
            plt.gca().xaxis.set_major_locator(ticker.MultipleLocator(243))
            plt.rcParams['font.sans-serif'] = ['SimHei']
            plt.show()

        elif isCSI300 and not isSSEC:
            ra = ResultAnalysis(initial=self.initial_money, start_date=start_date, end_date=end_date)
            days, return_list_CSI300 = ra.cal_return_list_CSI300()
            # num_days = len(return_list_SSEC)
            # days = list(range(num_days))
            # print(len(days))
            fig, ax = plt.subplots()
            ax.ticklabel_format(style='plain')
            # 根据情况画一个初始资金虚线
            plt.plot(days, return_list, color='red')
            plt.plot(days, return_list_CSI300, color='blue')
            plt.xlabel('时间')
            plt.ylabel('总资产（元）')
            plt.axhline(ra.get_initial_money(), linestyle='--', color='black')
            plt.gca().xaxis.set_major_locator(ticker.MultipleLocator(243))
            plt.rcParams['font.sans-serif'] = ['SimHei']
            plt.show()

        elif isSSEC and isCSI300:
            ra = ResultAnalysis(initial=self.initial_money, start_date=start_date, end_date=end_date)
            days, return_list_SSEC = ra.cal_return_list_SSEC()
            _, return_list_CSI300 = ra.cal_return_list_CSI300()

            fig, ax = plt.subplots()
            ax.ticklabel_format(style='plain')

            plt.plot(days, return_list, color='red')
            plt.plot(days, return_list_SSEC, color='green')
            plt.plot(days, return_list_CSI300, color='blue')
            plt.title('资产变化曲线图')
            plt.xlabel('时间')
            plt.ylabel('总资产（元）')
            plt.legend([strategy, '一直持有上证指数', '一直持有沪深300指数'])
            plt.axhline(ra.get_initial_money(), linestyle='--', color='black')
            plt.gca().xaxis.set_major_locator(ticker.MultipleLocator(243))
            plt.rcParams['font.sans-serif'] = ['SimHei']
            plt.show()

    def plot_gain(self, gain_list, isSSEC=False, isCSI300=False, start_date='2010-01-04', end_date='2019-12-31'):
        '''
        绘制收益变化曲线，投资开始和结束时间可以自由设定，可以不是交易日。
        :param gain_list: 策略收益变化曲线
        :param isSSEC: 是否绘制上证指数
        :param isCSI300: 是否绘制沪深300指数
        :param start_date: 策略开始时间
        :param end_date: 策略结束时间
        '''
        if isSSEC and not isCSI300:
            ra = ResultAnalysis(initial=self.initial_money, start_date=start_date, end_date=end_date)
            days, gain_list_SSEC = ra.cal_gain_list_SSEC()

            fig, ax = plt.figure()
            ax.ticklabel_format(style='plain')
            # 根据情况画一个初始资金虚线
            plt.plot(days, gain_list, color='red')
            plt.plot(days, gain_list_SSEC, color='green')
            plt.xlabel('时间')
            plt.ylabel('总收益（元）')
            plt.axhline(0, linestyle='--', color='black')
            plt.gca().xaxis.set_major_locator(ticker.MultipleLocator(243))
            plt.rcParams['font.sans-serif'] = ['SimHei']
            plt.show()

        elif isCSI300 and not isSSEC:
            ra = ResultAnalysis(initial=self.initial_money, start_date=start_date, end_date=end_date)
            days, gain_list_CSI300 = ra.cal_gain_list_CSI300()

            fig, ax = plt.figure()
            ax.ticklabel_format(style='plain')
            # 根据情况画一个初始资金虚线
            plt.plot(days, gain_list, color='red')
            plt.plot(days, gain_list_CSI300, color='blue')
            plt.xlabel('时间')
            plt.ylabel('总收益（元）')
            plt.axhline(0, linestyle='--', color='black')
            plt.gca().xaxis.set_major_locator(ticker.MultipleLocator(243))
            plt.rcParams['font.sans-serif'] = ['SimHei']
            plt.show()

        elif isSSEC and isCSI300:
            ra = ResultAnalysis(initial=self.initial_money, start_date=start_date, end_date=end_date)
            days, gain_list_SSEC = ra.cal_gain_list_SSEC()
            _, gain_list_CSI300 = ra.cal_gain_list_CSI300()

            fig, ax = plt.subplots()
            ax.ticklabel_format(style='plain')

            plt.plot(days, gain_list, color='red')
            plt.plot(days, gain_list_SSEC, color='green')
            plt.plot(days, gain_list_CSI300, color='blue')
            plt.xlabel('时间')
            plt.ylabel('总资产（元）')
            plt.legend(['策略', '一直持有上证指数', '一直持有沪深300指数'])
            plt.axhline(ra.get_initial_money(), linestyle='--', color='black')
            plt.gca().xaxis.set_major_locator(ticker.MultipleLocator(243))
            plt.rcParams['font.sans-serif'] = ['SimHei']
            plt.show()

    @staticmethod
    def group_bar_chart(labels, strategy_score, SSEC_score, strategy='策略'):
        x = np.arange(len(labels))
        width = 0.35
        opacity = 0.4

        fig, ax = plt.subplots()
        rects1 = ax.bar(x - width / 2, strategy_score, width,
                        alpha=opacity, color='b', label=strategy)
        rects2 = ax.bar(x + width / 2, SSEC_score, width,
                        alpha=opacity, color='r', label='一直持有上证指数')
        ax.set_ylabel('比率')
        ax.set_xlabel('评估指标')
        ax.set_title('各项指标对比图')
        ax.set_xticks(x)
        ax.set_xticklabels(labels)
        ax.legend()

        Visualization.autolabel(rects1, ax)
        Visualization.autolabel(rects2, ax)
        # autolabel(rects1)
        # autolabel(rects2)

        fig.tight_layout()

    @staticmethod
    def group_3bar_chart(labels, strategy1_score, strategy2_score, SSEC_score, strategy1='涨幅策略', strategy2='双均线策略'):
        x = np.arange(len(labels))
        x *= 2
        width = 0.35
        opacity = 0.4

        fig, ax = plt.subplots()
        rects1 = ax.bar(x - width, strategy1_score, width,
                        alpha=opacity, color='b', label=strategy1)
        rects2 = ax.bar(x, strategy2_score, width,
                        alpha=opacity, color='r', label=strategy2)
        rects3 = ax.bar(x + width, SSEC_score, width,
                        alpha=opacity, color='g', label='一直持有上证指数')

        ax.set_ylabel('比率')
        ax.set_xlabel('评估指标')
        ax.set_title('各项指标对比图')
        ax.set_xticks(x)
        ax.set_xticklabels(labels)
        ax.legend()

        Visualization.autolabel(rects1, ax)
        Visualization.autolabel(rects2, ax)
        Visualization.autolabel(rects3, ax)
        # autolabel(rects1)
        # autolabel(rects2)

        fig.tight_layout()

    @staticmethod
    def autolabel(rects, ax):
        """Attach a text label above each bar in *rects*, displaying its height."""
        for rect in rects:
            height = rect.get_height()
            ax.annotate('{}'.format(height),
                        xy=(rect.get_x() + rect.get_width() / 2, height),
                        xytext=(0, 3),  # 3 points vertical offset
                        textcoords="offset points",
                        ha='center', va='bottom')
